//! Flight SQL客户端实现

use std::sync::Arc;
use tonic::transport::Channel;
use arrow_flight::{
    flight_service_client::FlightServiceClient,
    FlightDescriptor, Ticket,
};
use arrow::record_batch::RecordBatch;
use data_gateway_kernel::Result;

/// Flight SQL客户端配置
#[derive(Debug, Clone)]
pub struct FlightSqlClientConfig {
    /// 服务器地址
    pub server_address: String,
    /// 连接超时时间（秒）
    pub connect_timeout_seconds: u64,
    /// 请求超时时间（秒）
    pub request_timeout_seconds: u64,
}

/// Flight SQL客户端
pub struct FlightSqlClient {
    /// 配置
    config: FlightSqlClientConfig,
    /// gRPC客户端
    client: FlightServiceClient<Channel>,
}

impl Default for FlightSqlClientConfig {
    fn default() -> Self {
        Self {
            server_address: "http://127.0.0.1:50051".to_string(),
            connect_timeout_seconds: 30,
            request_timeout_seconds: 300,
        }
    }
}

impl FlightSqlClient {
    /// 创建新的Flight SQL客户端
    pub async fn new(config: FlightSqlClientConfig) -> Result<Self> {
        tracing::info!("连接Flight SQL服务器: {}", config.server_address);
        
        let channel = Channel::from_shared(config.server_address.clone())
            .map_err(|e| anyhow::anyhow!("无效的服务器地址: {}", e))?
            .connect_timeout(std::time::Duration::from_secs(config.connect_timeout_seconds))
            .timeout(std::time::Duration::from_secs(config.request_timeout_seconds))
            .connect()
            .await
            .map_err(|e| anyhow::anyhow!("连接服务器失败: {}", e))?;
        
        let client = FlightServiceClient::new(channel);
        
        tracing::info!("Flight SQL客户端连接成功");
        
        Ok(Self {
            config,
            client,
        })
    }
    
    /// 执行SQL查询
    pub async fn execute_query(&mut self, sql: &str) -> Result<RecordBatch> {
        tracing::debug!("执行SQL查询: {}", sql);
        
        // 创建Flight描述符
        let descriptor = FlightDescriptor {
            r#type: arrow_flight::flight_descriptor::DescriptorType::Cmd as i32,
            cmd: Some(sql.as_bytes().to_vec()),
            path: vec![],
        };
        
        // 获取Flight信息
        let flight_info = self.client
            .get_flight_info(descriptor)
            .await
            .map_err(|e| anyhow::anyhow!("获取Flight信息失败: {}", e))?
            .into_inner();
        
        // 创建Ticket
        let ticket = Ticket {
            ticket: sql.as_bytes().to_vec(),
        };
        
        // 执行查询
        let mut stream = self.client
            .do_get(ticket)
            .await
            .map_err(|e| anyhow::anyhow!("执行查询失败: {}", e))?
            .into_inner();
        
        // 读取结果
        let mut record_batches = Vec::new();
        let mut schema = None;
        
        while let Some(flight_data) = stream.message().await
            .map_err(|e| anyhow::anyhow!("读取查询结果失败: {}", e))? {
            
            if schema.is_none() {
                // 解析Schema
                let message = arrow::ipc::root_as_message(&flight_data.data_header)
                    .map_err(|e| anyhow::anyhow!("解析Schema消息失败: {}", e))?;
                
                if message.header_type() == arrow::ipc::MessageHeader::Schema {
                    let ipc_schema = message.header_as_schema().unwrap();
                    schema = Some(arrow::ipc::convert::fb_to_schema(ipc_schema));
                }
            } else {
                // 解析数据
                let message = arrow::ipc::root_as_message(&flight_data.data_header)
                    .map_err(|e| anyhow::anyhow!("解析数据消息失败: {}", e))?;
                
                if message.header_type() == arrow::ipc::MessageHeader::RecordBatch {
                    let record_batch = arrow::ipc::reader::read_record_batch(
                        &flight_data.data_body,
                        message,
                        schema.as_ref().unwrap(),
                        &arrow::ipc::reader::IpcReadOptions::default(),
                        None,
                    ).map_err(|e| anyhow::anyhow!("解析RecordBatch失败: {}", e))?;
                    
                    record_batches.push(record_batch);
                }
            }
        }
        
        // 合并结果
        if record_batches.is_empty() {
            Ok(RecordBatch::new_empty(Arc::new(
                schema.unwrap_or_else(|| arrow::datatypes::Schema::empty())
            )))
        } else if record_batches.len() == 1 {
            Ok(record_batches.into_iter().next().unwrap())
        } else {
            // 合并多个RecordBatch
            arrow::compute::concat_batches(
                &record_batches[0].schema(),
                &record_batches,
            ).map_err(|e| anyhow::anyhow!("合并RecordBatch失败: {}", e).into())
        }
    }
    
    /// 获取服务器信息
    pub async fn get_server_info(&mut self) -> Result<Vec<String>> {
        let action = arrow_flight::Action {
            r#type: "get_server_info".to_string(),
            body: vec![],
        };
        
        let mut stream = self.client
            .do_action(action)
            .await
            .map_err(|e| anyhow::anyhow!("获取服务器信息失败: {}", e))?
            .into_inner();
        
        let mut info = Vec::new();
        while let Some(result) = stream.message().await
            .map_err(|e| anyhow::anyhow!("读取服务器信息失败: {}", e))? {
            
            let info_str = String::from_utf8_lossy(&result.body);
            info.push(info_str.to_string());
        }
        
        Ok(info)
    }
    
    /// 获取配置
    pub fn config(&self) -> &FlightSqlClientConfig {
        &self.config
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_flight_sql_client_config_default() {
        let config = FlightSqlClientConfig::default();
        assert_eq!(config.server_address, "http://127.0.0.1:50051");
        assert_eq!(config.connect_timeout_seconds, 30);
        assert_eq!(config.request_timeout_seconds, 300);
    }
    
    #[test]
    fn test_flight_sql_client_config_custom() {
        let config = FlightSqlClientConfig {
            server_address: "http://localhost:8080".to_string(),
            connect_timeout_seconds: 60,
            request_timeout_seconds: 600,
        };
        
        assert_eq!(config.server_address, "http://localhost:8080");
        assert_eq!(config.connect_timeout_seconds, 60);
        assert_eq!(config.request_timeout_seconds, 600);
    }
}
